# -*- coding: utf-8 -*-
import os

import torchvision
import torchvision.datasets.mnist
import tb

# torchvision.datasets.CIFAR10(root="./data", download=True) # CIFAR10数据集

with tb.tb() as writer:
    for i, img in enumerate(os.listdir("./data/structure1/train/bees")):
        writer.add_image("bees", torchvision.io.read_image(f"./data/structure1/train/bees/{img}"), i)

''' dataloader '''
from torch.utils.data import DataLoader
''' batch_size : batch的大小
    shuffle : 是否打乱数据
    num_workers : 读取数据的线程数
    pin_memory : 是否将数据拷贝到GPU的内存中
    drop_last : 是否丢掉最后一个不完整的batch
    collate_fn : 默认为None, 表示使用默认的collate_fn, 否则使用自定义的collate_fn
'''
dataset = torchvision.datasets.MNIST(root="./data/MNIST", train=False)
loader = DataLoader(dataset=dataset, batch_size=4, shuffle=True, num_workers=0, drop_last=False) # 一个简单例子